# -*- coding: utf-8 -*-

# File Name： resnet_block
# Description :
# Author : lirui
# create_date： 2022/6/3
# Change Activity:
from torch import nn
from torch.nn import Conv2d

from .... import is_same_shape, logger
from ..base.base_module import BaseModule
from ..layers.norm import build_norm_layer
from ...misc.util import int_to_tuple
from ....cv.modules.resize_tensor import resize_tensor


class ResnetBlock(BaseModule):
    def __init__(self, in_channels=32,
                 out_channels=32,
                 kernel_size=3,
                 stride=1,
                 dilation=1,
                 expansion=1,
                 norm_cfg=None,
                 ):
        """
        Args:
            in_channels (int): Number of channels in the input image
            out_channels (int): Number of channels produced by the convolution
            kernel_size (int or tuple): Size of the convolving kernel
            stride (int or tuple, optional): Stride of the convolution. Default: 1
        """
        super().__init__()

        if norm_cfg is None:  # this module must contain norm layer.
            norm_cfg = dict(type='BN')
        self.expansion = expansion
        norm_cfg['num_features'] = out_channels * self.expansion
        self.norm_cfg = norm_cfg
        self.norm1 = build_norm_layer(norm_cfg)
        self.norm2 = build_norm_layer(norm_cfg)
        self.conv1 = Conv2d(in_channels=in_channels,
                            out_channels=out_channels * self.expansion,
                            kernel_size=int_to_tuple(kernel_size),
                            stride=int_to_tuple(stride),
                            padding=dilation,
                            dilation=int_to_tuple(dilation))
        self.conv2 = Conv2d(in_channels=out_channels * self.expansion,
                            out_channels=out_channels * self.expansion,
                            kernel_size=int_to_tuple(kernel_size),
                            padding=1,
                            bias=False)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = self._build_downsample(in_channels, out_channels, stride)
        self.stride = stride
        self.dilation = dilation

    def _build_downsample(self, in_channels, out_channels, stride):
        """

        Args:
            in_channels:
            out_channels:
            stride:

        Returns:

        """
        if stride != 1 or in_channels != out_channels * self.expansion:
            return nn.Sequential(
                Conv2d(in_channels=in_channels, out_channels=out_channels * self.expansion, kernel_size=int_to_tuple(stride), stride=int_to_tuple(stride)),
                build_norm_layer(self.norm_cfg)
            )
        return None

    def forward(self, x):
        """
        Forward function.
        """
        identity = x.clone()
        out = self.conv1(x)
        out = self.norm1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.norm2(out)
        if self.downsample is not None:
            identity = self.downsample(identity)
        if not is_same_shape(out, identity):
            logger.warning('Identity and out are not same size.')
            identity = resize_tensor(identity, out.shape[-2:])
        out += identity
        out = self.relu(out)
        return out
